//
// Created by DonAzufre on 24-5-19.
//

#include "llvm/Pass.h"
#include "llvm/IR/Function.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/IR/IRBuilder.h"

using namespace llvm;

namespace {
    using namespace std;

    struct ShowBranchTF : public ModulePass {
        static char ID;

        ShowBranchTF() : ModulePass(ID) {}

        bool runOnModule(Module &M) override {
            auto &ctx = M.getContext();
            auto printfArgTy = PointerType::getUnqual(ctx);
            auto printfFuncTy = FunctionType::get(IntegerType::getInt32Ty(ctx), printfArgTy, true);
            auto printfCallee = M.getOrInsertFunction("printf", printfFuncTy);

            // 使用LLVM19编译时, 以下四行取消注释后可以正常编译加载
            // 使用LLVM14, legacy passmanager时, 如果加上以下四行, 会发生segment fault
            // 目前没有找到原因

            // PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
            // Stack dump:
            // 0.      Program arguments: ../llvm-project-14/build/bin/opt -load build/ShowBranchTF/ShowBranchTFPass.so -show-branch-tf -enable-new-pm=0 test/fib.ll
            // 1.      Running pass 'Show true or false of each condbr at runtime.' on module 'test/fib.ll'.
            // #0 0x000055a4dac26434 PrintStackTraceSignalHandler(void*) Signals.cpp:0:0
            // #1 0x000055a4dac23bdb SignalHandler(int) Signals.cpp:0:0
            // #2 0x00007f622b05b050 (/lib/x86_64-linux-gnu/libc.so.6+0x3c050)
            // #3 0x000055a4da294831 llvm::Function::addFnAttr(llvm::Attribute::AttrKind) (../llvm-project-14/build/bin/opt+0x25f5831)
            // #4 0x00007f622b45cbb5 llvm::Function::setDoesNotThrow() /home/iris/projects/llvm-project-14/llvm/include/llvm/IR/Function.h:561:3
            // #5 0x00007f622b45b91f (anonymous namespace)::ShowBranchTF::runOnModule(llvm::Module&) /home/iris/projects/llvm-passes/ShowBranchTF/ShowBranchTF.cpp:29:33
            // #6 0x000055a4da2ed658 llvm::legacy::PassManagerImpl::run(llvm::Module&) (../llvm-project-14/build/bin/opt+0x264e658)
            // #7 0x000055a4d85367c3 main (../llvm-project-14/build/bin/opt+0x8977c3)
            // #8 0x00007f622b04624a (/lib/x86_64-linux-gnu/libc.so.6+0x2724a)
            // #9 0x00007f622b046305 __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x27305)
            // #10 0x000055a4d8609241 _start (../llvm-project-14/build/bin/opt+0x96a241)
            // [1]    4609 segmentation fault (core dumped)  ../llvm-project-14/build/bin/opt -load build/ShowBranchTF/ShowBranchTFPass.so


//            auto printf = dyn_cast<Function>(printfCallee.getCallee());
//            printf->setDoesNotThrow();
//            printf->addParamAttr(0, llvm::Attribute::NoCapture);
//            printf->addParamAttr(0, llvm::Attribute::ReadOnly);

            auto trueString = ConstantDataArray::getString(
                    ctx, "[ShowBranchTF - condbr %d] Branch True.\n"
            );
            auto falseString = ConstantDataArray::getString(
                    ctx, "[ShowBranchTF - condbr %d] Branch False.\n"
            );

            auto tsVar = M.getOrInsertGlobal("trueString", trueString->getType());
            auto fsVar = M.getOrInsertGlobal("falseString", falseString->getType());

            dyn_cast<GlobalVariable>(tsVar)->setInitializer(trueString);
            dyn_cast<GlobalVariable>(fsVar)->setInitializer(falseString);


            for (auto &func: M) {
                if (func.isDeclaration())
                    continue;

                int condbrNum = 0;

                for (auto &blk: func) {
                    for (auto &inst: blk) {
                        if (auto br = dyn_cast<BranchInst>(&inst); br != nullptr && br->isConditional()) {
                            // br已经是带条件的分支了, 一定有两个后继
                            assert(br->getNumSuccessors() == 2);
                            auto trueSucc = br->getSuccessor(0);
                            auto falseSucc = br->getSuccessor(1);

                            IRBuilder<> builder(&*trueSucc->getFirstInsertionPt());
                            auto tsPtr = builder.CreatePointerCast(tsVar, printfArgTy, "tsptr");
                            builder.CreateCall(printfCallee, {tsPtr, builder.getInt32(condbrNum)});

                            builder.SetInsertPoint(&*falseSucc->getFirstInsertionPt());
                            auto fsPtr = builder.CreatePointerCast(fsVar, printfArgTy, "fsptr");
                            builder.CreateCall(printfCallee, {fsPtr, builder.getInt32(condbrNum)});

                            condbrNum++;
                        }
                    }
                }
            }

            return false;
        }
    };
}

char ShowBranchTF::ID = 0;
static RegisterPass<ShowBranchTF> X("show-branch-tf", "Show true or false of each condbr at runtime.");